
#include "libmri.h"

#include <stdarg.h>
#include <stdlib.h>
#include <math.h>

#include <sys/time.h>

#include <algorithm>
#include <iostream>


using std::max;
using std::min;
using std::cerr;
using std::endl;
using std::cout;


/* calib3d_chol() implements the same operation as calib3d_posv(): find a
 * least-squares least-norm solution to the calibration equations. However,
 * it leverages the similarity in the n_channels matrices to solve each of the
 * calibration equations by cholesky-factorizing the original calib3d_matrix()
 */
void calib3d_chol (cplxf *&kernels, int *ksize, cplxf *calib, int *csize, float lambda)
{
	int cs1 = csize[0], cs2 = csize[1], cs3 = csize[2], nch = csize[3],
	    ks1 = ksize[0], ks2 = ksize[1], ks3 = ksize[2];

	int n =	ks1*ks2*ks3*nch;
	int m1 = cs1-ks1+1, m2 = cs2-ks2+1, m3 = cs3-ks3+1, m = m1*m2*m3;

	kernels = new cplxf [n*nch];

	for (int i = 0; i < n*nch; ++i){ kernels[i] = 0.0f; }

	complex *A = new complex [m*n],
	      *AtA = new complex [n*n],
	        *L = new complex [n*n];

	calib3d_matrix ((cplxf*) A, calib, csize, ksize);

	complex _1 = {1.0f, 0.0f}, _0 = {0.0f, 0.0f}, _n1 = {-1.0f, 0.0f};

	cgemm ('C', 'N', n, n, m, &_1, A, m, A, m, &_0, AtA, n);

	double frob_AtA = 0;

	for (int i = 0; i < n*n; ++i){
		frob_AtA += ((cplxf*) AtA)[i].abs2();
		L[i] = AtA[i];
	}
	frob_AtA = sqrt (frob_AtA);

	for (int i = 0; i < n; ++i)
		L[i + n*i].real += frob_AtA*lambda/n;

	int info;

	cpotrf ('L', n, L, n, &info);

	complex *bt = new complex [n];
	complex *bh = new complex[n];
	complex *e = new complex[n];
	complex *eh = new complex[n];
	complex *Bh = new complex [2*n];
	complex *Eh = new complex [2*n];
	complex *t0 = new complex [2*n];
	complex *t1 = new complex [n*n];
	complex *x = new complex [n];

	for (int ch = 0; ch < nch; ++ch)
	{
		int iy = (ks1/2) + ks1*((ks2/2) + ks2*((ks3/2) + ks3*ch));

		for (int i = 0; i < n; ++i)
			bt[i] = AtA[i + n*iy];

		for (int i = 0; i < n; ++i)
			bh[i] = bt[i];
		ctrsv ('L', 'N', 'N', n, L, n, bh, 1);

		for (int i = 0; i < n; ++i)
			if (i == iy) e[i] = compose_complex (1.0f,0.0f);
			        else e[i] = compose_complex (0.0f,0.0f);

		for (int i = 0; i < n; ++i)
			eh[i] = e[i];
		ctrsv ('L', 'N', 'N', n, L, n, eh, 1);

		for (int i = 0; i < n; ++i){
			Bh[i] = complex_negate (bh[i]);
			Bh[i+n] = complex_negate (eh[i]);

			Eh[i] = eh[i];
			Eh[i+n] = bh[i];
		}

		complex T[4] = {{1.0f,0.0f}, {0.0f,0.0f},
		                {0.0f,0.0f}, {1.0f,0.0f}};
		int Tpiv[2];

		cgemm ('C', 'N', 2, 2, n, &_1, Eh, n, Bh, n, &_1, T, 2);
		cgetrf (2, 2, T, 2, Tpiv, &info);
		cgetri (2, T, 2, Tpiv, &info);

		cgemm ('N', 'C', 2, n, 2, &_1, T, 2, Eh, n, &_0, t0, 2); 

		for (int j = 0; j < n; ++j)
			for (int i = 0; i < n; ++i)
				if (i == j) t1[i + n*j] = compose_complex(1.0f,0.0f);
				       else t1[i + n*j] = compose_complex(0.0f,0.0f);

		cgemm ('N', 'N', n, n, 2, &_n1, Bh, n, t0, 2, &_1, t1, n);
		cgemv ('N', n, n, &_1, t1, n, bh, 1, &_0, x, 1);
		ctrsv ('L', 'C', 'N', n, L, n, x, 1);

		x[iy] = compose_complex (0.0f, 0.0f);

		for (int i = 0; i < n; ++i)
			kernels [i + n*ch] = x[i];
	}

	delete [] A; delete [] AtA; delete [] L;
	delete [] bt; delete [] bh;
	delete [] e; delete [] eh;
	delete [] Bh; delete [] Eh; 
	delete [] t0; delete [] t1;
	delete [] x;
}

#ifdef _MAIN_

int main (int argc, char **argv)
{
	shm_array_t *shm_calib = 0;
	int *ksize = 0, nksize = 0;
	float lambda = .02;

	struct arg* args [] = {
		new list_arg_t<int> ("ksize", ksize, nksize, "x"),
		new arg_t<shm_array_t*> ("calib", shm_calib),
		new arg_t<float> ("lambda", lambda),
		0
	};

	get_args (argc, argv, args);

	if (!ksize){
		ksize = new int[3];
		ksize[0] = ksize[1] = ksize[2] = 7;
	}

	if (shm_calib == 0){
		if (get_shm_array("calib", &shm_calib)){
			perror ("calib");
			return -1;
		}
	}
	if (shm_calib->type != shm_array_cplxf){
		cerr << "calib must be single-precision complex!" << endl;
		return -1;
	}
	if (shm_calib->ndims != 4){
		cerr << "calib must be <int>x<int>x<int>x<int>" << endl;
		return -1;
	}
	cplxf *calib = (cplxf *) shm_calib->data.c;

	int csize[4] = {shm_calib->dims[0], shm_calib->dims[1],
	                shm_calib->dims[2], shm_calib->dims[3]};

	shm_array_t shm_kernels_size, *shm_kernels = 0;

	shm_kernels_size.type = shm_array_cplxf;
	shm_kernels_size.ndims = 5;
	for (int i = 0; i < 3; ++i)
		shm_kernels_size.dims[i] = ksize[i];
	shm_kernels_size.dims[3] = shm_kernels_size.dims[4] = csize[3];

	if (make_shm_array ("kernels", &shm_kernels_size, &shm_kernels)){
		perror("kernels");
		return -1;
	}
	cplxf *kernels = (cplxf *) shm_kernels->data.c;

	cout << "ksize = " << ksize[0] << " x " << ksize[1] << " x " << ksize[2] <<
	endl << "csize = " << csize[0] << " x " << csize[1] << " x " << csize[2] <<
	endl << "coils = " << csize[3] << ", lambda = " << lambda << endl;


	calib3d_chol (kernels, calib, csize, ksize, lambda);

	return 0;
}

#endif
